###############################################################
###############################################################
#####
#####		Chase Joyner
#####	       882 Homework 3
#####
###############################################################
###############################################################

###############
## Problem 1 ##

## Part c ##
normal.invgamma = function(Y, X, mu0, R0, a, b){
	
	## Preliminaries ##
	n = length(Y)
	p = dim(X)[2]
	iter = 1e4
	
	## Save Records ##
	Beta = matrix(-99, ncol = iter, nrow = p)
	Sigmasq = rep(-99, iter)

	## Initial values ##
	beta = rep(1, p)
	sigmasq = 1

	## Calculations ##
	IR = solve(t(X) %*% X + solve(R0))
	mu = IR %*% (t(X) %*% Y + solve(R0) %*% mu0)

	for(i in 1:iter){
		## Update Beta ##
		beta = t(rmvnorm(1, mu, sigmasq[1] * IR))
		
		## Update Sigmasq ##
		sigmasq = rinvgamma(1,(n+p+a)/2, (1/2)*(t(Y-X%*%beta)%*%(Y-X%*%beta)+t(beta-mu0)%*%IR%*%(beta-mu0)+b))

		## Save records ##
		Beta[,i] = beta
		Sigmasq[i] = sigmasq
	}
	
	return(list(Beta = Beta, Sigmasq = Sigmasq))
}


## Part d ##
library(MCMCpack)
library(mvtnorm)
library(coda)
beta.true = c(0.5,1,3)
sigmasq.true = 1
n = 1000
eta = 500	## historical sample size
p = length(beta.true)
X = cbind(rep(1,n), rbinom(n, 1, 0.5), rnorm(n, 2, 1))
alpha0 = 1
g = 0.1
sims = 100
Beta.power = matrix(-99, ncol = sims, nrow = p)
Sigmasq.power = rep(-99, sims)
Beta.g = matrix(-99, ncol = sims, nrow = p)
Sigmasq.g = rep(-99, sims)
Beta.mlr = matrix(-99, ncol = sims, nrow = p)
Sigmasq.mlr = rep(-99, sims)
## Power initial priors ##
mu0 = rep(2,p)
R = diag(100, p)
## G initial priors ##
beta0 = rep(2,p)
a = b = 1
for(i in 1:sims){
	Y = rnorm(n, X %*% beta.true, sqrt(sigmasq.true))

	## For power prior ##
	#Y0 = rnorm(eta, X %*% beta.true, sqrt(sigmasq.true))
	#X0 = cbind(rep(1,eta), rbinom(eta, 1, 0.5), rnorm(eta, 2, 1))
	Y0 = Y
	X0 = X
	C.power = solve(t(X0)%*%X0 + solve(alpha0*R))
	mu.power = C.power %*% (t(X0)%*%Y0 + solve(alpha0*R)%*%mu0)
	cov.power = sigmasq.true / alpha0 * C.power
	alpha.power = eta*alpha0
	beta.power = alpha0*(t(Y0)%*%Y0-t(mu.power)%*%solve(C.power)%*%mu.power+t(mu0)%*%solve(alpha0*R)%*%mu0)
	res.power = normal.invgamma(Y,X,mu.power,cov.power,alpha.power,beta.power)
	Beta.power[,i] = apply(res.power$Beta,1,mean)
	Sigmasq.power[i] = mean(res.power$Sigmasq)

	## For g prior ##
	mu.g = beta0
	cov.g = g*sigmasq.true*solve(t(X)%*%X)
	alpha.g = a
	beta.g = b
	res.g = normal.invgamma(Y,X,mu.g,cov.g,alpha.g,beta.g)
	Beta.g[,i] = apply(res.g$Beta,1,mean)
	Sigmasq.g[i] = mean(res.power$Sigmasq)

	## Multiple Linear Regression in class ##
	res.mlr = normal.invgamma(Y,X,rep(2,p),sigmasq.true*R,1,1)
	Beta.mlr[,i] = apply(res.mlr$Beta,1,mean)
	Sigmasq.mlr[i] = mean(res.mlr$Sigmasq)

	print(i)
}

apply(Beta.power,1,mean)
mean(Sigmasq.power)
apply(Beta.g,1,mean)
mean(Sigmasq.g)
apply(Beta.mlr,1,mean)
mean(Sigmasq.mlr)





###############
## Problem 2 ##

## Part c ##
full.gibbs = function(Y, D, W, rho, b, sig2, tau2, as2, bs2, at2, bt2, iter = 1e4, verbose = TRUE){
	n = length(Y)
	DW = D - rho * W
	I = sparseMatrix(1:n, 1:n, x = 1)
	#I = diag(n)

	b.save = matrix(-99, nrow = n, ncol = iter)
	sig2.save = rep(-99, n)
	tau2.save = rep(-99, n)
	beta0.save = rep(-99, n)

	for(t in 1:iter){
		## Sample intercept ##
		beta0 = rnorm(1, mean(Y - b), sqrt(sig2 / n))

		## Sample spatial random effects all at once ##
		Prec = I + DW / tau2
		CH = chol(Prec)
		R1 = solve(CH, rnorm(n, 0, sqrt(sig2)), sparse = TRUE)
		R2 = solve(t(CH), Y - beta0, sparse = TRUE)
		R3 = solve(CH, R2, sparse = TRUE)
		b = as.vector(R1 + R3)

		## Sample sig2 ##
		as2s = n + as2
		bs2s = as.vector(bs2 + (1/2) * (sum((Y - beta0 - b)^2) + b %*% DW %*% b / tau2))
		sig2 = 1 / rgamma(1, as2s, bs2s)

		## Sample tau2 ##
		at2s = n / 2 + at2
		bt2s = bt2 + (1/2) * (t(b) %*% DW %*% b / sig2)
		tau2 = 1 / rgamma(1, at2s, bt2s)

		## Save ##
		beta0.save[t] = beta0
		b.save[,t] = b
		sig2.save[t] = sig2
		tau2.save[t] = tau2

		if(verbose == TRUE){
			print(t)
			if(t %% 100 == 0){
				par(mfrow = c(3,1))
				plot(beta0.save[1:t])
				plot(sig2.save[1:t])
				plot(tau2.save[1:t])
			}
		}
	}
	return(list("sig2" = sig2.save, "tau2" = tau2.save, "beta0" = beta0.save, "b" = b.save))
}


one.gibbs = function(Y, D, W, rho, b, sig2, tau2, as2, bs2, at2, bt2, iter = 1e4, verbose = TRUE){
	n = length(Y)
	DW = D - rho * W
	D = diag(D)

	b.save = matrix(-99, nrow = n, ncol = iter)
	sig2.save = rep(-99, n)
	tau2.save = rep(-99, n)
	beta0.save = rep(-99, n)

	for(t in 1:iter){
		## Sample intercept ##
		beta0 = rnorm(1, mean(Y - b), sqrt(sig2 / n))

		## Sample spatial random effects one at a time ##
		for(i in 1:n){
			mui = rho * W[,i] %*% b / D[i]
			Wtau = D[i] / tau2
			pm = as.vector((Y[i] - beta0 + Wtau * mui) / (1 + Wtau))
			ps = as.vector(sqrt(sig2 / (1 + Wtau)))
			b[i] = rnorm(1, pm, ps)
		}

		## Sample sig2 ##
		as2s = n + as2
		bs2s = as.vector(bs2 + (1/2) * (sum((Y - beta0 - b)^2) + t(b) %*% DW %*% b / tau2))
		sig2 = 1 / rgamma(1, as2s, bs2s)

		## Sample tau2 ##
		at2s = n / 2 + at2
		bt2s = bt2 + (1/2) * (t(b) %*% DW %*% b / sig2)
		tau2 = 1 / rgamma(1, at2s, bt2s)

		## Save ##
		beta0.save[t] = beta0
		b.save[,t] = b
		sig2.save[t] = sig2
		tau2.save[t] = tau2

		if(verbose == TRUE){
			print(t)
			if(t %% 100 == 0){
				par(mfrow = c(3,1))
				plot(beta0.save[1:t])
				plot(sig2.save[1:t])
				plot(tau2.save[1:t])
			}
		}
	}
	return(list("sig2" = sig2.save, "tau2" = tau2.save, "beta0" = beta0.save, "b" = b.save))
}


## Part d ##
library(lattice)
library(MASS)
library(Matrix) 
library(mvtnorm)
library(hierarchicalDS)
P1 = 50
P2 = 50
Y = matrix(-99, P1, P2)
for(p1 in 1:P1){
	for(p2 in 1:P2){
		Y[p1,p2] = 10*dnorm((p1-p2), 0, 10) + rnorm(1, 0, 0.1)
	}
}
levelplot(Y, main = "Observed Data")
W = square_adj(P1)
D = diag(rowSums(W))
Y = as.vector(Y)
n = length(Y)

load(file.choose())


res1 = full.gibbs(Y=Y,D=D,W=W,rho=1,b=rep(0,n),sig2=1,tau2=1,as2=0.001,bs2=0.001,at2=0.001,bt2=0.001)
res2 = one.gibbs(Y=Y,D=D,W=W,rho=1,b=rep(0,n),sig2=1,tau2=1,as2=0.001,bs2=0.001,at2=0.001,bt2=0.001)

path = "C:/Users/Norm/Desktop/882/HW 3/"
load(paste0(path, "res1.RData"))
load(paste0(path, "res2.RData"))

Yhat1 = mean(res1$beta0[5000:10000]) + apply(res1$b[,5000:10000], 1, mean)
levelplot(matrix(Yhat1, ncol = P1, nrow = P2), main = "Full Block Gibbs")
par(mfrow = c(3,1))
plot(res1$beta0[5000:10000]+res1$b[1,5000:10000], ylab = "beta0 + b1", main = "Full block gibbs")
plot(res1$sig2[5000:10000], ylab = "sig2")
plot(res1$tau2[5000:10000], ylab = "tau2")

Yhat2 = mean(res2$beta0[5000:10000]) + apply(res2$b[,5000:10000], 1, mean)
levelplot(matrix(Yhat2, ncol = P1, nrow = P2), main = "Individual Gibbs")
par(mfrow = c(3,1))
plot(res2$beta0[5000:10000]+res2$b[1,5000:10000], ylab = "beta0 + b1", main = "Individual gibbs")
plot(res2$sig2[5000:10000], ylab = "sig2")
plot(res2$tau2[5000:10000], ylab = "tau2")